#pragma once

#include <vector>
#include <opencv2/opencv.hpp>
#include "yolo_define.h"

_YOLO_NAMESPACE_BEGINE



    class Box
    {
    public:
        float left, top, right, bottom, confidence;
        int class_label;

        Box() = default;
        Box(const Box &b) = default;
        Box(float left, float top, float right, float bottom, float confidence, int class_label)
            : left(left), top(top), right(right), bottom(bottom), confidence(confidence), class_label(class_label) {}

        float area() const;

        Box operator|(const Box &b) const;

        Box operator&(const Box &b) const;

        bool operator==(const Box &b) const
        {
            return left == b.left && top == b.top && right == b.right && bottom == b.bottom;
        }

        float iou(const Box &b) const;

        friend std::ostream& operator<<(std::ostream& os, const Box& box)
        {
            os << "Box(" << box.left << ", " << box.top << ", " << box.right << ", " << box.bottom << ", " << box.confidence << ", " << box.class_label << ")";
            return os;
        }
    };

    // float iou(const Box &a, const Box &b)
    // {
    //     return a.iou(b);
    // }

    typedef std::vector<Box> BoxArray;
    static BoxArray nms_cpu(BoxArray &boxes, float nms_threshold)
    {
        std::sort(boxes.begin(), boxes.end(), [](const Box &a, const Box &b)
                  { return a.confidence > b.confidence; });
        BoxArray res;
        res.reserve(boxes.size());
        std::vector<bool> remove_flag(boxes.size(), false);
        for (int i = 0; i < boxes.size(); ++i)
        {
            if (remove_flag[i])
                continue;
            auto curBox = boxes[i];
            res.emplace_back(curBox);
            for (int j = i + 1; j < boxes.size(); ++j)
            {
                if (remove_flag[j])
                    continue;
                auto iou = curBox.iou(boxes[j]);
                if (iou > nms_threshold && boxes[j].class_label == curBox.class_label)
                    remove_flag[j] = true;
            }
        }
        return res;
    }

    struct AffineMatrix
    {
        float i2d[6]; // image to dst(network), 2x3 matrix
        float d2i[6]; // dst to image, 2x3 matrix

        void compute(const cv::Size &from, const cv::Size &to)
        {
            float scale_x = to.width / (float)from.width;
            float scale_y = to.height / (float)from.height;
            float scale = std::min(scale_x, scale_y);
            i2d[0] = scale;
            i2d[1] = 0;
            i2d[2] = -scale * from.width * 0.5 + to.width * 0.5 + scale * 0.5 - 0.5;
            i2d[3] = 0;
            i2d[4] = scale;
            i2d[5] = -scale * from.height * 0.5 + to.height * 0.5 + scale * 0.5 - 0.5;

            cv::Mat m2x3_i2d(2, 3, CV_32F, i2d);
            cv::Mat m2x3_d2i(2, 3, CV_32F, d2i);
            cv::invertAffineTransform(m2x3_i2d, m2x3_d2i);
        }

        cv::Mat i2d_mat()
        {
            return cv::Mat(2, 3, CV_32F, i2d);
        }
    };

_YOLO_NAMESPACE_END
